% this version does not loop over sims to construct the 
% path 
% note this only allows for CARA utility 
% for CRRA need to change the PoliFunctions script first

function [C,A] = C_Savings(Wi,MU,CAPH,gamma,iX,rho,r,amax)

% rho is discount rate (beta)
% r is the interest rate 

Nstates = size(CAPH,1);  

% dropping the "dead" state
% makes the future sum to be taken only 
% over the alive states
% which means that the individual puts 0 value to the dead 
% state, consistent with how we are writting the lifetime utility function

n = Nstates-1;

MUs = squeeze(MU);
MUs = MUs(1:n,:)';
Ygrid = repmat(Wi,1,n)-MUs;
Ygrid = permute(Ygrid, [3 2 1]);        % Permute dimensions to make it consistent with original code
% put in dollars
Ygrid = Ygrid*1000;

% Size of the grid of shocks

T = size(CAPH,3)+1;

%Ygrid = repmat(ygrid,1,1,T);

Abar   = 0;                          % Borrowing constraint (min level of assets)
agrid3 = linspace(Abar,amax*1000,101);

[AT2, CT2]=PoliFunctions(Ygrid,agrid3,T,gamma,rho,r,CAPH(1:n,1:n,:));

% Simulate Paths

[Rows, Cols] = size(iX);

sims = Cols;


    yhistt2 = zeros(T,sims);
    ahistt2 = zeros(T,sims);
    dhistt2 = zeros(T,sims);
    chistt2 = zeros(T,sims);
    
    iXt = iX(1,:);
    iXt(iXt==n+1)=1;
    ahistt2(1,:) = AT2(iXt,1);
    yhistt2(1,:) = Ygrid(1,iXt,1);
    chistt2(1,:) = CT2(iXt,1);
    dhistt2(1,:) = zeros(1,sims);
 
    for i = 2:T
    
        % a_ind is 1xsims that gives the column number 
        % for the asset in the previous period
 
    a= repmat(ahistt2(i-1,:)',1,101)==repmat(agrid3,sims,1);
    a_ind = sum(a.*repmat(1:101,sims,1),2)';

        %iXt gives the row number
    iXt = iX(i,:);
    dhistt2(i,:) = iXt==(n+1);
    % replace the index for 1 if index is     
    iXt(iXt==n+1)=1;
    yhistt2(i,:) = Ygrid(1,iXt,i).*(1-dhistt2(i,:));
    AT2t = AT2(:,:,i);
    CT2t = CT2(:,:,i);
    % linear index
    index = n*(a_ind-1)+iXt;
    ahistt2(i,:) = AT2t(index).*(1-dhistt2(i,:));
    chistt2(i,:) = CT2t(index).*(1-dhistt2(i,:));
    
    end

%YF2 = [YF2;yhistt2];
%CF2 = [CF2;chistt2];
%AF2 = [AF2;ahistt2];
%DF2 = [DF2;dhistt2];

%end

C = chistt2;
A = ahistt2;

end
